# model.py

import torch.nn as nn
from torchcrf import CRF

# 在model.py文件中修改LstmCRFModel类
class LstmCRFModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=1, dropout=0.1):
        super(LstmCRFModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.output_seq_length = output_seq_length
        
        # 添加dropout参数
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True, dropout=dropout if num_layers > 1 else 0)
        self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
        self.dropout = nn.Dropout(dropout)  # 添加dropout层
        self.crf = CRF(output_dim, batch_first=True)
        
    def forward(self, x, target=None, mask=None):
        # LSTM层
        lstm_out, _ = self.lstm(x)
        
        # 应用dropout
        lstm_out = self.dropout(lstm_out)
        
        # 全连接层
        fc_out = self.fc(lstm_out[:, -1, :])
        fc_out = fc_out.view(-1, self.output_seq_length, self.output_dim)
        
        # 如果提供了目标和掩码，计算损失
        if target is not None and mask is not None:
            loss = -self.crf(fc_out, target, mask=mask, reduction='mean')
            return loss
        
        # 否则进行解码
        predictions = self.crf.decode(fc_out)
        return predictions
